package frege.test.HashSets where

import frege.data.List (sort)
import Data.Bits (.&., shiftR)

data HashSet a = Empty 
               | ! Branch { b0, b1, b2, b3, b4, b5, b6, b7, b8, b9, bA, bB, bC, bD, bE, bF :: HashSet a }
               | ! Bucket [a]
               
derive Eq HashSet a   

instance Show Show a => HashSet a where
    show s = "fromList " ++ show (toList s)            

instance Empty HashSet where   
  empty = Empty
  null Empty = true
  null _ = false

instance ListSource HashSet where                
  toList s = go s [] where
     go Empty ys = ys
     go (Bucket xs) ys = xs ++ ys
     go (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) ys = 
       go x0 (go x1 (go x2 (go x3 (go x4 (go x5 (go x6 (go x7 (go x8 
       (go x9 (go xA (go xB (go xC (go xD (go xE (go xF ys)))))))))))))))

class SetOp s where
  union :: s -> s -> s
  difference :: s -> s -> s
  intersection :: s -> s -> s  
  exclusion :: s -> s -> s
  
instance SetOp (Eq a) => HashSet a where
  union Empty s2 = s2
  union s1 Empty = s1
  union (Bucket xs1) (Bucket xs2) = Bucket (xs1 ++ filter (\x2 -> not $ elem x2 xs1) xs2)
  union branch1 branch2 = zipBranch union branch1 branch2   

  difference Empty _ = Empty
  difference s1 Empty = s1
  difference (Bucket xs1) (Bucket xs2) = emptyOrBucket $ filter (\x1 -> not $ elem x1 xs2) xs1 
  difference branch1 branch2 = zipBranch difference branch1 branch2

  intersection :: Eq a => HashSet a -> HashSet a -> HashSet a 
  intersection Empty _ = Empty
  intersection _ Empty = Empty          
  intersection (Bucket xs1) (Bucket xs2) = emptyOrBucket $ filter (\x1 -> elem x1 xs2) xs1 
  intersection branch1 branch2 = zipBranch intersection branch1 branch2

  exclusion Empty s2 = s2
  exclusion s1 Empty = s1
  exclusion (Bucket xs1) (Bucket xs2) = emptyOrBucket $
     filter (\x1 -> not $ elem x1 xs2) xs1 ++
     filter (\x2 -> not $ elem x2 xs1) xs2
  exclusion branch1 branch2 = zipBranch exclusion branch1 branch2   

unions :: SetOp s => [s] -> s
unions ss = foldr1 union ss    

intersections :: SetOp s => [s] -> s
intersections ss = foldr1 intersection ss    

singleton :: Eq a => a -> HashSet a
singleton x = insert x empty


insert :: Eq a => a -> HashSet a -> HashSet a
insert x s = ins (hashCode x) 0 s where
   ins _ 7 Empty = Bucket [x]
   ins _ 7 (bucket @ Bucket xs) = if Prelude.elem x xs then bucket else Bucket (x:xs)
   ins h k Empty = changeBranch emptyBranch (h `.&.` 15) (const $ ins (h `shiftR` 4) (k + 1) Empty)
   ins h k branch = changeBranch branch (h `.&.` 15) $ ins (h `shiftR` 4) (k + 1)

delete :: Eq a => a -> HashSet a -> HashSet a
delete x s = del (hashCode x) s where
   del _ Empty = Empty
   del _ (Bucket xs) = emptyOrBucket $ filter (!= x) xs
   del h branch = replaceEmptyBranch $ changeBranch branch (h `.&.` 15) $ del (h `shiftR` 4)
   replaceEmptyBranch branch | branch == emptyBranch = Empty
                             | otherwise = branch
       
member :: Eq a => a -> HashSet a -> Bool
member x s = mbr s (hashCode x) where
   mbr Empty _ = false
   mbr (Bucket xs) _ = elem x xs  
   mbr branch h = mbr (getBranch branch (h `.&.` 15)) (h `shiftR` 4)
         
partition :: Eq a => (a -> Bool) -> HashSet a -> (HashSet a, HashSet a)
partition f s = foldr insPair (empty, empty) $ toList s where
    insPair x (left, right) | f x = (insert x left, right)
                            | otherwise = (left, insert x right)            
          
fromList :: Eq a => [a] -> HashSet a         
fromList xs = fold ins empty xs where
   ins s x = insert x s 

size :: HashSet a -> Int
size Empty = 0
size (Bucket xs) = length xs
size (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) =
  size x0 + size x1 + size x2 + size x3 + size x4 + size x5 + size x6 + size x7 + 
  size x8 + size x9 + size xA + size xB + size xC + size xD + size xE + size xF     

toOrdList :: Ord a => HashSet a -> [a]
toOrdList Empty = []
toOrdList (Bucket xs) = sort xs
toOrdList (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) =
   (((tm x0 x1) `m` (tm x2 x3)) `m` ((tm x4 x5) `m` (tm x6 x7))) `m`
   (((tm x8 x9) `m` (tm xA xB)) `m` ((tm xC xD) `m` (tm xE xF))) where
     tm set1 set2 = toOrdList set1 `m` toOrdList set2                                
     m [] ys = ys
     m xs [] = xs
     m (xxs @ x:xs) (yys @ y:ys) | x < y = x : m xs yys
                                 | otherwise = y : m xxs ys
                
private emptyBranch :: HashSet a
private emptyBranch = Branch Empty Empty Empty Empty Empty Empty Empty Empty
                             Empty Empty Empty Empty Empty Empty Empty Empty

private changeBranch :: HashSet a -> Int -> (HashSet a -> HashSet a) -> HashSet a
private changeBranch br  0 f = br.{ b0 <- f } 
private changeBranch br  1 f = br.{ b1 <- f } 
private changeBranch br  2 f = br.{ b2 <- f } 
private changeBranch br  3 f = br.{ b3 <- f } 
private changeBranch br  4 f = br.{ b4 <- f } 
private changeBranch br  5 f = br.{ b5 <- f } 
private changeBranch br  6 f = br.{ b6 <- f } 
private changeBranch br  7 f = br.{ b7 <- f } 
private changeBranch br  8 f = br.{ b8 <- f } 
private changeBranch br  9 f = br.{ b9 <- f } 
private changeBranch br 10 f = br.{ bA <- f } 
private changeBranch br 11 f = br.{ bB <- f } 
private changeBranch br 12 f = br.{ bC <- f } 
private changeBranch br 13 f = br.{ bD <- f } 
private changeBranch br 14 f = br.{ bE <- f } 
private changeBranch br  _ f = br.{ bF <- f } 

private getBranch :: HashSet a -> Int -> HashSet a
private getBranch (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) n = case n of 
   0 -> x0; 1 -> x1; 2 -> x2; 3 -> x3; 4 -> x4; 5 -> x5; 6 -> x6; 7 -> x7;
   8 -> x8; 9 -> x9; 10 -> xA; 11 -> xB; 12 -> xC; 13 -> xD; 14 -> xE; _ -> xF
   
private zipBranch :: (HashSet a -> HashSet a -> HashSet a) -> HashSet a -> HashSet a -> HashSet a   
private zipBranch f (Branch x0 x1 x2 x3 x4 x5 x6 x7 x8 x9 xA xB xC xD xE xF) 
                    (Branch y0 y1 y2 y3 y4 y5 y6 y7 y8 y9 yA yB yC yD yE yF) = 
   Branch (f x0 y0) (f x1 y1) (f x2 y2) (f x3 y3) (f x4 y4) (f x5 y5) (f x6 y6) (f x7 y7)     
          (f x8 y8) (f x9 y9) (f xA yA) (f xB yB) (f xC yC) (f xD yD) (f xE yE) (f xF yF)   

private emptyOrBucket :: [a] -> HashSet a
private emptyOrBucket [] = Empty
private emptyOrBucket xs = Bucket xs 


   
     
                
     